# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jhkim

import datetime
import os
import re
import pandas
from absl import app, flags
import logging

import coin.research.read_strat_info as rsi
import coin.research.flag_generator as fgen
import coin.strategy.mm.tool.archive_base as abase
from python.coin_deploy.lm_agg2 import from_coin1_generator as fromcoin1gen
import hashlib


def format_cmd(execcmd, flags_dict, td_is_str=False):
  cmd = " ".join([f"--{key}={value}" for key, value in flags_dict.items() if value is not None])
  cmd = f"{execcmd} {cmd}"
  if flags.FLAGS.extra_args is not None:
    cmd = f"{cmd} {flags.FLAGS.extra_args}"

  if flags.FLAGS.use_feed_cache:
    return cmd
  else:
    return f"RUST_FEED_HANDLER=1 RUST_TOPIC_RECORD=1 {cmd}"


def fmt_values(flags_dict):
  out_str = '_'.join([val.replace("/", "_") for val in flags_dict.values() if val is not None])
  return int(hashlib.sha1(out_str.encode('utf-8')).hexdigest(), 16) % (10**8)


def main(_):
  if flags.FLAGS.slurm_cpu == '0':
    flags.FLAGS.slurm_cpu = None
  cmds = []
  for cmd in flags.FLAGS.cmd.split(','):
    flags.FLAGS.cmd = cmd
    cmds.extend(launch())
  execute_cmds(cmds)


def launch():
  cmd_relpath = os.path.relpath(os.path.abspath(__file__))
  nparallel = flags.FLAGS.nparallel

  if cmd_relpath.startswith("python/"):
    root_path = "python"
  else:
    root_path = "."

  tds = abase.get_trading_dates(flags.FLAGS.trading_date)

  if flags.FLAGS.subsymbols is not None:
    fs_overrides = list(abase.get_flags_with_full_symbols(flags.FLAGS.subsymbols))

  cmds = []
  if flags.FLAGS.cmd == 'pta':
    strat_opts = []
    if flags.FLAGS.strategy_group is not None:
      strat_opts.append(
          (f" --strat_group {flags.FLAGS.strategy_group}", flags.FLAGS.strategy_group))
    elif flags.FLAGS.strategy_name is not None:
      stratnames = []
      stratinfos = rsi.get_strategy_info(trading_date=tds,
                                         strategy_group=flags.FLAGS.strategy_group,
                                         strategy_name=flags.FLAGS.strategy_name,
                                         subsymbols=flags.FLAGS.subsymbols)
      stratnames.extend([stratinfo['strategy_name'] for stratinfo in stratinfos])

      for stratname in sorted(list(set(stratnames))):
        strat_opts.append((f"--strat_name {stratname}", stratname))
    else:
      raise ValueError("strategy_group or strategy_name should be supported")
    for strat_opt, prefix in strat_opts:
      cmds.append(f"./pyrunner {root_path}/coin/support/latency/calc_taker_ratio_one_day.py"
                  + f" --date {flags.FLAGS.trading_date} {strat_opt} --prefix {prefix}_"
                  + f" && ./pyrunner {root_path}/coin/support/latency/demo_tr.py"
                  + f" --date {flags.FLAGS.trading_date} --prefix {prefix}_")
  elif flags.FLAGS.cmd == 'pnl_plot':
    stratnames = []
    stratinfos = rsi.get_strategy_info(trading_date=tds,
                                       strategy_group=flags.FLAGS.strategy_group,
                                       strategy_name=flags.FLAGS.strategy_name,
                                       subsymbols=flags.FLAGS.subsymbols,
                                       pick_symbol_regex=flags.FLAGS.pick_symbol_regex)
    stratdf = pandas.DataFrame(stratinfos)
    grpdf = stratdf.groupby(
        list(set(stratdf.columns)
        - set(['trading_date', 'symbol', 'subsymbol'])))
    if 'subsymbol' in grpdf:
      stratdf = grpdf.nunique()['subsymbol'].reset_index()
    for _, stratinfo in stratdf.iterrows():
      try:
        if flags.FLAGS.strategy_name is not None and not re.match(flags.FLAGS.strategy_name, stratinfo['strategy_name']):
          continue
      except:
        continue
      flags_dict = {}
      if flags.FLAGS.dir is not None:
        if flags.FLAGS.fold:
          flags_dict['out_dir'] = os.path.join(flags.FLAGS.dir, stratinfo['strategy_name'])
        else:
          flags_dict['out_dir'] = flags.FLAGS.dir
      if flags.FLAGS.time_range is not None:
        flags_dict['time_range'] = flags.FLAGS.time_range
      if flags.FLAGS.pick_symbol_regex is not None:
        flags_dict['pick_symbol_regex'] = f'"{flags.FLAGS.pick_symbol_regex}"'
      if flags.FLAGS.split_debug_hours is not None:
        flags_dict['split_debug_hours'] = flags.FLAGS.split_debug_hours
        if flags.FLAGS.full_bbo:
          flags_dict['full_bbo'] = True
      if flags.FLAGS.timed_ret is not None:
        flags_dict['timed_ret'] = flags.FLAGS.timed_ret
      if flags.FLAGS.slurm_cpu is not None:
        cpu_cnt = 1  # max(int(stratinfo['subsymbol'] / 2), 1)
        cmd_prefix = f"srun -N 1 -c {cpu_cnt} -n {flags.FLAGS.slurm_cpu} --exclusive " +\
            f" --job-name=pplot_{stratinfo['strategy_name']}"
      else:
        cmd_prefix = ""
      cmd_format = (f"{cmd_prefix} ./pyrunner {root_path}/coin/strategy/mm/tool/pnl_plot.py"
                    + f" --trading_date {flags.FLAGS.trading_date} "
                    + f" --feed_machine feed-05.ap-northeast-1.aws"
                    + f" --orderlog_machine {stratinfo['machine']}"
                    + f" --strategy_name {stratinfo['strategy_name']}")
      if (len(abase.get_trading_dates(flags.FLAGS.trading_date)) >= 30 and
          fromcoin1gen.is_team_sci_strat_group_fcs_mt(stratinfo['strategy_group'])):
        cmd_format += " --is_focus_mt=1"
      cmds.append(format_cmd(cmd_format, flags_dict))
  elif flags.FLAGS.cmd in ['sim']:
    cmds.append("mkdir -p log_sim")
    stratinfos = rsi.get_strategy_info(trading_date=tds[0],
                                       strategy_group=flags.FLAGS.strategy_group,
                                       strategy_name=flags.FLAGS.strategy_name,
                                       subsymbols=flags.FLAGS.subsymbols,
                                       pick_symbol_regex=flags.FLAGS.pick_symbol_regex)
    name_conf_map = rsi.get_strat_name_conf_map(root_path)
    td = flags.FLAGS.trading_date
    for stratinfo in stratinfos:
      flags_dict = {}
      flags_dict['trading_date'] = td
      if stratinfo['strategy_name'] not in name_conf_map:
        print("# omitting %s" % stratinfo['strategy_name'])
        continue
      outdict = name_conf_map[stratinfo['strategy_name']]
      stratfile, sim_type = outdict['stratfile'], outdict['sim_type']
      flags_dict['sim_type'] = sim_type
      if flags.FLAGS.force_rerun:
        flags_dict['force_rerun'] = 'True'
      if flags.FLAGS.cmd == 'sim':
        flags_dict['config_path'] = stratfile
        if flags.FLAGS.sparams is not None:
          flags_dict['sparams'] = f'\"{flags.FLAGS.sparams}\"'
        #if flags.FLAGS.slurm_cpu is not None:
        #  flags_dict['slurm_cpu'] = flags.FLAGS.slurm_cpu
        if flags.FLAGS.cpu is not None:
          flags_dict['cpu'] = flags.FLAGS.cpu
        if flags.FLAGS.split_debug_hours is not None:
          flags_dict['split_debug_hours'] = flags.FLAGS.split_debug_hours
        if flags.FLAGS.full_bbo:
          flags_dict['full_bbo'] = True
        if flags.FLAGS.hours is not None:
          flags_dict['hours'] = flags.FLAGS.hours
        else:
          time_range = flags.FLAGS.time_range
          if time_range is not None:
            time_range = time_range.replace('-', 'H-') + 'H'
          flags_dict['hours'] = time_range
        if flags.FLAGS.dir is not None:
          if flags.FLAGS.fold:
            flags_dict['sim_result_dir'] = os.path.join(flags.FLAGS.dir, stratinfo['strategy_name'])
          else:
            flags_dict['sim_result_dir'] = flags.FLAGS.dir
        if stratinfo['strategy_group'].find("lm") >= 0 or stratinfo['strategy_group'] in [
            'interval'
        ]:
          flags_dict['sim_type'] = 'lm'
        else:
          flags_dict['sim_type'] = 'basis_smm'
        cmds.append(
            format_cmd(f"./pyrunner cc/appcoin2/strategy/run_sim.py", flags_dict, td_is_str=True))
        if flags.FLAGS.slurm_cpu is not None:
          cmds[-1] = f"srun -N 1 -c 1 -n {flags.FLAGS.slurm_cpu} --exclusive " +\
              f" --job-name=ru_{stratinfo['strategy_name']} {cmds[-1]}"
      else:
        raise ValueError(flags.FLAGS.cmd)
  elif flags.FLAGS.cmd in ['order_plot', 'kuona', 'replay', 'order_accept_latency']:
    stratinfos = rsi.get_strategy_info(trading_date=tds[0],
                                       strategy_group=flags.FLAGS.strategy_group,
                                       strategy_name=flags.FLAGS.strategy_name,
                                       subsymbols=flags.FLAGS.subsymbols,
                                       pick_symbol_regex=flags.FLAGS.pick_symbol_regex)

    log_dir = os.path.join(flags.FLAGS.dir or '.', "log")
    os.makedirs(log_dir, exist_ok=True)
    cmds.append(f"mkdir -p {log_dir}")
    for td in [flags.FLAGS.trading_date]:
      for i, stratinfo in enumerate(stratinfos):
        flags_dict = {}

        if flags.FLAGS.cmd in ['order_plot', 'order_accept_latency']:
          if flags.FLAGS.dir is not None:
            flags_dict['out_dir'] = os.path.join(flags.FLAGS.dir, stratinfo['strategy_name'])
          if flags.FLAGS.cmd in ['order_plot']:
            flags_dict['title'] = stratinfo['strategy_name'] + "_"

        flags_dict['trading_date'] = td
        flags_dict['time_range'] = fgen.get_time_range(flags.FLAGS.trading_date,
                                                       flags.FLAGS.time_range)
        flags_dict['feed_machine'] = fgen.get_feed_machine(order_machine=stratinfo['machine'])
        flags_dict['orderlog_machine'] = stratinfo['machine']
        if flags.FLAGS.cmd == 'order_accept_latency':
          flags_dict['mea'] = f"{stratinfo['market_type']}:{stratinfo['exchange']}:"
        else:
          flags_dict['subsymbols'] = \
              f"{stratinfo['exchange']}:{stratinfo['market_type']}:{stratinfo['subsymbol']}"
        if flags.FLAGS.force_rerun:
          flags_dict['force_rerun'] = 'True'
        if flags.FLAGS.api_override is not None:
          flags_dict['api_override'] = flags.FLAGS.api_override
        if flags.FLAGS.full_bbo:
          flags_dict['full_bbo'] = True
        if flags.FLAGS.owner is not None:
          flags_dict['owner'] = flags.FLAGS.owner
        else:
          flags_dict['owner'] = stratinfo['owner']

        if flags.FLAGS.slurm_cpu is not None:
          cmd_prefix = f"srun -n {flags.FLAGS.slurm_cpu} --mem-per-cpu 2g --job-name=pplot_{stratinfo['strategy_name']}"
        else:
          cmd_prefix = ""
        cmds.append(format_cmd(
            f"{cmd_prefix} ./pyrunner {root_path}/coin/strategy/mm/tool/{flags.FLAGS.cmd}.py",
            flags_dict,
            td_is_str=True)
            + f" &> {log_dir}/{fmt_values(flags_dict)} & # {stratinfo['strategy_name']}")
  elif flags.FLAGS.cmd == 'feed_plot':
    for fs_override in fs_overrides:
      flags_dict = {}
      flags_dict['trading_date'] = flags.FLAGS.trading_date
      flags_dict['time_range'] = fgen.get_time_range(flags.FLAGS.trading_date,
                                                     flags.FLAGS.time_range)
      flags_dict['feed_machine'] = fgen.get_feed_machine(exchange=fs_override['exchange'])
      flags_dict.update(fs_override)
      cmds.append(
          format_cmd(f"./pyrunner {root_path}/coin/strategy/mm/tool/feed_archive_checker.py",
                     flags_dict))

      flags_dict["compare_worker_ids"] = "1-2-1,2"
      cmds.append(
          format_cmd(f"./pyrunner {root_path}/coin/strategy/mm/tool/feed_archive_checker.py",
                     flags_dict))

      del flags_dict["compare_worker_ids"]
      flags_dict["compare_feed_machines"] =\
          "feed-01.ap-northeast-1.aws,feed-01.cn-hongkong.aliyun" \
          + ",feed-01.eu-west-1.aws,feed-01.us-west-1.aws"
      cmds.append(
          format_cmd(f"./pyrunner {root_path}/coin/strategy/mm/tool/feed_archive_checker.py",
                     flags_dict))
  else:
    raise ValueError(flags.FLAGS.cmd)
  return cmds


def execute_cmds(cmds):
  cmd_out = ""
  cmds = list(set(cmds))
  if (flags.FLAGS.slurm_cpu is not None
      and flags.FLAGS.nparallel > 0):
    tsstr = f'{datetime.datetime.now().strftime("%Y%m%d_%H%M%S_%f")}'
    job_filename = f'/tmp/sbatch_{tsstr}.sh'
    sbatch_out_file = os.path.realpath(f'sbatch_out.{tsstr}.txt')
    with open(job_filename, 'w') as job_file:
      wrapcmds = (
          ["#!/bin/bash"]
          + [f"#SBATCH -o {sbatch_out_file} "]
          + [""]
          + [f"{cmd} &" for cmd in cmds] + ["wait"])
      jobstr = "\n".join(wrapcmds)
      job_file.write(jobstr)
    print("# of jobs: %s" % len(cmds))
    print(job_filename)
    print(sbatch_out_file)
    nbatch = min(len(cmds), flags.FLAGS.nparallel)
    job_name = flags.FLAGS.cmd.replace(",", "_")
    cmd_out = f"sbatch -n {nbatch} --job-name={job_name} -c 1 --mem-per-cpu=4gb --wait {job_filename}"
    print(cmd_out)
    if flags.FLAGS.run:
      os.system(cmd_out)
  else:
    ucmds = {}
    for cmd in cmds:
      ucmds[cmd] = True
    for i, cmd in enumerate(list(ucmds.keys())):
      cmd_out += cmd + " &" + "\n"
      if (i + 1) % flags.FLAGS.nparallel == 0:
        cmd_out += "wait & \n"
    print(cmd_out)
    if flags.FLAGS.run:
      from cc.appcoin2.strategy.sim_driver import batch_local
      batch_local(cmds, flags.FLAGS.nparallel)


#  bash experimental/jhkim/pta_lmagg.sh --trading_date=20191223-20191230
#  ./pyrunner coin/research/cmd_generator.py --strategy_group=lm-agg \
#  --cmd=sim --trading_date=20191223-20191230 --latency_multiplier=1,2
if __name__ == '__main__':
  flags.DEFINE_string('strategy_group', None, '')
  flags.DEFINE_string('strategy_name', None, '')
  # cmd : order_plot
  flags.DEFINE_string('cmd', None, '')
  flags.DEFINE_string('subsymbols', None, '')
  flags.DEFINE_string('latency_multiplier', None, '')
  flags.DEFINE_integer('nparallel', 5, '')
  flags.DEFINE_integer('timed_ret', None, '')
  flags.DEFINE_string('sparams', None, '')
  flags.DEFINE_string('slurm_cpu', None, '')
  flags.DEFINE_string('cpu', None, '')
  flags.DEFINE_string('hours', None, '')
  flags.DEFINE_string('dir', None, '')
  flags.DEFINE_string('pick_symbol_regex', None, '')
  flags.DEFINE_float('split_debug_hours', None, '')
  flags.DEFINE_bool('full_bbo', None, '')
  flags.DEFINE_bool('assert_feed_cache', None, '')
  flags.DEFINE_bool('force_rerun', None, '')
  flags.DEFINE_bool('fold', False, '')
  flags.DEFINE_bool('run', False, '')
  flags.DEFINE_string('extra_args', None, '')

  abase.define_base_flags()
  abase.define_feed_archive_flags()

  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
